use image::{DynamicImage, GenericImageView};
use ndarray::{Array, ArrayBase, Axis, Dim, Ix, OwnedRepr, s};
use ort::SessionOutputs;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct BoundingBox {
    pub x1: f32,
    pub y1: f32,
    pub x2: f32,
    pub y2: f32,
}

pub fn intersection(box1: &BoundingBox, box2: &BoundingBox) -> f32 {
    (box1.x2.min(box2.x2) - box1.x1.max(box2.x1)) * (box1.y2.min(box2.y2) - box1.y1.max(box2.y1))
}

pub fn union(box1: &BoundingBox, box2: &BoundingBox) -> f32 {
    ((box1.x2 - box1.x1) * (box1.y2 - box1.y1)) + ((box2.x2 - box2.x1) * (box2.y2 - box2.y1))
        - intersection(box1, box2)
}

pub fn pixels_to_array(img: DynamicImage) -> ArrayBase<OwnedRepr<f32>, Dim<[Ix; 4]>> {
    let shape: Dim<[Ix; 4]> = Dim([1, 3, img.width() as _, img.height() as _]);
    let mut input: ArrayBase<OwnedRepr<f32>, Dim<[Ix; 4]>> = Array::zeros(shape);
    for pixel in img.pixels() {
        let x = pixel.0 as _;
        let y = pixel.1 as _;
        let [r, g, b, _] = pixel.2.0;
        input[[0, 0, y, x]] = (r as f32) / 255.;
        input[[0, 1, y, x]] = (g as f32) / 255.;
        input[[0, 2, y, x]] = (b as f32) / 255.;
    }
    input
}

pub fn convert_result<'a>(
    outputs: &SessionOutputs,
    names: &'a Vec<&str>,
    original_img: &DynamicImage,
) -> eyre::Result<Vec<(BoundingBox, &'a str, f32)>> {
    let (img_width, img_height) = (original_img.width(), original_img.height());

    let output = outputs["output0"]
        .try_extract_tensor::<f32>()?
        .t()
        .into_owned();

    let mut boxes = Vec::new();
    let output = output.slice(s![.., .., 0]);
    for row in output.axis_iter(Axis(0)) {
        let row: Vec<_> = row.iter().copied().collect();
        let (class_id, prob) = row
            .iter()
            // skip bounding box coordinates
            .skip(4)
            .enumerate()
            .map(|(index, value)| (index, *value))
            .reduce(|accum, e| if e.1 > accum.1 { e } else { accum })
            .unwrap();
        if prob < 0.5 {
            continue;
        }
        let label = names[class_id];
        let xc = row[0] / 640. * (img_width as f32);
        let yc = row[1] / 640. * (img_height as f32);
        let w = row[2] / 640. * (img_width as f32);
        let h = row[3] / 640. * (img_height as f32);
        boxes.push((
            BoundingBox {
                x1: xc - w / 2.,
                y1: yc - h / 2.,
                x2: xc + w / 2.,
                y2: yc + h / 2.,
            },
            label,
            prob,
        ));
    }

    boxes.sort_by(|box1, box2| box2.2.total_cmp(&box1.2));
    let mut result = Vec::new();

    while !boxes.is_empty() {
        result.push(boxes[0]);
        boxes = boxes
            .iter()
            .filter(|box1| intersection(&boxes[0].0, &box1.0) / union(&boxes[0].0, &box1.0) < 0.7)
            .copied()
            .collect();
    }

    Ok(result)
}

